package chapter10

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.expressions.Window

/**
 * author: yuhui
 * descriptions: 用dsl风格api实现sql中的窗口分析函数
 * date: 2024 - 12 - 02 4:39 下午
 */
object DML_DSLAPI_WINDOW {

  def main(args: Array[String]): Unit = {

    val spark: SparkSession = SparkSession.builder()
      .appName(this.getClass.getSimpleName)
      .master("local[*]")
      .getOrCreate()

    val df = spark.read.option("header", true).csv("BookData/input/10stu2.csv")

    import spark.implicits._
    import org.apache.spark.sql.functions._

    val window = Window.partitionBy(Symbol("city")).orderBy(Symbol("score").desc)

    df.select(Symbol("id"), Symbol("name"), Symbol("age"), Symbol("sex"), Symbol("city"), Symbol("score"), row_number().over(window) as "rn")
      .where(Symbol("rn") <= 2)
      .drop("rn") // 最后结果中不需要rn列，可以drop掉这个列
      .select("id", "name", "age", "sex", "city", "score") // 或者用select指定你所需要的列
      .show()

    spark.close()
  }

}